import os
from typing import Dict, List, Any, Optional
from zhipuai import ZhipuAI
import re

# 导入math_a.py中的函数
from math_a import (
    # 基础数学函数
    base_add, base_subtract, base_multiply, base_divide,
    base_mod, base_floor_divide,
    
    # 三角数学函数
    trig_sin, trig_cos, trig_tan, trig_asin, trig_acos, trig_atan,
    trig_csc, trig_sec, trig_cot,
    
    # 高等数学函数
    adv_exp, adv_log, adv_pow, adv_sqrt,
    adv_log10, adv_log2, adv_factorial
)

class MathAgent:
    def __init__(self):
        self.client = ZhipuAI(api_key=os.getenv("GLM_API_KEY"))
        self.model = "glm-4-plus"
        self.searchModel = "Search-Std"
        self.conversation_history = []
        self.results = {}  # 存储步骤结果的字典
        
        # 准备函数库
        self.base_functions = {
            "base_add": base_add,
            "base_subtract": base_subtract,
            "base_multiply": base_multiply,
            "base_divide": base_divide,
            "base_mod": base_mod,
            "base_floor_divide": base_floor_divide
        }
        
        self.trig_functions = {
            "trig_sin": trig_sin,
            "trig_cos": trig_cos,
            "trig_tan": trig_tan,
            "trig_asin": trig_asin,
            "trig_acos": trig_acos,
            "trig_atan": trig_atan,
            "trig_csc": trig_csc,
            "trig_sec": trig_sec,
            "trig_cot": trig_cot
        }
        
        self.adv_functions = {
            "adv_exp": adv_exp,
            "adv_log": adv_log,
            "adv_pow": adv_pow,
            "adv_sqrt": adv_sqrt,
            "adv_log10": adv_log10,
            "adv_log2": adv_log2,
            "adv_factorial": adv_factorial
        }
        
        # 创建工具
        self.base_tool = self._create_tool(self.base_functions, "基础数学")
        self.trig_tool = self._create_tool(self.trig_functions, "三角数学")
        self.adv_tool = self._create_tool(self.adv_functions, "高等数学")
        
        # 简单表达式模式匹配
        self.simple_expr_patterns = [
            (r'([+-]?\d*\.?\d+)\s*\*\s*([+-]?\d*\.?\d+)', 'base_multiply', ['a', 'b']),
            (r'([+-]?\d*\.?\d+)\s*\+\s*([+-]?\d*\.?\d+)', 'base_add', ['a', 'b']),
            (r'([+-]?\d*\.?\d+)\s*-\s*([+-]?\d*\.?\d+)', 'base_subtract', ['a', 'b']),
            (r'([+-]?\d*\.?\d+)\s*/\s*([+-]?\d*\.?\d+)', 'base_divide', ['a', 'b']),
            (r'([+-]?\d*\.?\d+)\s*%\s*([+-]?\d*\.?\d+)', 'base_mod', ['a', 'b']),
            (r'([+-]?\d*\.?\d+)\s*//\s*([+-]?\d*\.?\d+)', 'base_floor_divide', ['a', 'b']),
        ]
        
        # 步骤结果引用模式
        self.result_ref_pattern = re.compile(r'\[步骤\s*(\d+)\s*结果\]')
    
    def _create_tool(self, functions: Dict[str, callable], category: str) -> Dict[str, Any]:
        """创建function scheme的工具"""
        tools = []
        for name, func in functions.items():
            # 获取函数的参数和文档字符串
            import inspect
            sig = inspect.signature(func)
            params = sig.parameters
            
            # 构建参数描述
            parameters = {}
            for param_name, param in params.items():
                param_info = {
                    "type": "number",
                    "description": f"参数 {param_name}"
                }
                parameters[param_name] = param_info
            
            # 构建函数描述
            description = func.__doc__ or f"{category}函数: {name}"
            
            # 添加到tools列表
            tools.append({
                "type": "function",
                "function": {
                    "name": name,
                    "description": description,
                    "parameters": {
                        "type": "object",
                        "properties": parameters,
                        "required": list(params.keys())
                    }
                }
            })
        
        return {
            "category": category,
            "tools": tools
        }
    
    def _determine_math_category(self, question: str) -> str:
        """确定数学问题的类别"""
        prompt = f"""
        请分析以下数学问题，是属于单步就可以处理的问题，还是需要多步骤处理的问题。属于哪个类别：
        问题: {question}
        
        可选类别:
        1. 基础数学 - 包括基本算术运算（加减乘除、取模、整除等）
        2. 三角数学 - 包括三角函数及其反函数、余割、正割、余切等
        3. 高等数学 - 包括指数、对数、幂运算、平方根、阶乘等
        
        请返回最匹配的类别名称，不要添加任何其他说明。
        """
        
        response = self.client.chat.completions.create(
            model=self.model,  # 使用GLM-4+模型进行分类,
            messages=[{"role": "user", "content": prompt}]
        )
        
        category = response.choices[0].message.content.strip()
        
        # 验证返回的类别是否有效
        valid_categories = ["基础数学", "三角数学", "高等数学"]
        if category not in valid_categories:
            # 如果无法确定，默认使用基础数学
            return "基础数学"
        
        return category
    
    def _execute_function_call(self, function_name: str, arguments: Dict[str, Any]) -> Any:
        """执行函数调用"""
        # 根据函数名选择正确的函数库
        if function_name.startswith("base_"):
            func = self.base_functions[function_name]
        elif function_name.startswith("trig_"):
            func = self.trig_functions[function_name]
        elif function_name.startswith("adv_"):
            func = self.adv_functions[function_name]
        else:
            raise ValueError(f"未知的函数名: {function_name}")
        
        # 执行函数
        try:
            # print(f"正在调用函数: {function_name}，参数: {arguments}")
            result = func(**arguments)
            print(f"函数 {function_name}, 参数: {arguments} 执行完成，结果: {result}")
            return result
        except Exception as e:
            return f"计算错误: {str(e)}"
    
    def _decompose_steps(self, question: str) -> List[Dict[str, Any]]:
        """分解问题为多个计算步骤"""
        # 尝试直接识别简单表达式
        if self._is_simple_expression(question):
            return [{"description": question, "result": None}]
            
        prompt = f"""
        请将以下数学问题分解为多个计算步骤，并按照顺序排列：
        问题: {question}
        
        请确保每个步骤都是明确的计算操作，不要包含解释性文字。
        每个步骤都以数字开头，以半角;结束。整个回复以半角;结束。
        如果某个步骤要引用上面某步骤的计算结果，请用[步骤X结果]表示，其中X为步骤序号。
        例如：  1.计算a+b;2.计算c*d;3.计算[步骤 1 结果]/[步骤 2 结果];...  
        """
        
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}]
        )
        
        steps_text = response.choices[0].message.content.strip()
        steps = []
        
        # 解析步骤文本
        # 更灵活的步骤解析，移除固定前缀限制
        import re
        step_pattern = re.compile(r'^(\d+)\.\s*(.*)$')
        for line in steps_text.split(';'):
            match = step_pattern.match(line.strip())
            if match:
                step_description = match.group(2).strip()
                steps.append({"description": step_description, "result": None})
                step_number = match.group(1)
                print(f"步骤 {step_number}: 分解步骤: {step_description}")
        return steps
    
    def _is_simple_expression(self, question: str) -> bool:
        """判断是否为简单表达式"""
        for pattern, _, _ in self.simple_expr_patterns:
            if re.search(pattern, question):
                return True
        return False
    
    def _parse_simple_expression(self, step_description: str, steps: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        """解析简单表达式为函数调用，替换步骤结果引用"""
        # 先替换步骤结果引用
        step_description = self._replace_result_references(step_description, steps)
        print(f'替换结果引用后的步骤描述: {step_description}')

        for pattern, function_name, param_names in self.simple_expr_patterns:
            match = re.search(pattern, step_description)
            if match:
                args = match.groups()
                # print(f'匹配到的参数组: {args}')
                if len(args) == len(param_names):
                    arguments = {param_names[i]: float(args[i]) for i in range(len(args))}
                    # print(f'生成的参数: {arguments}')
                    return {
                        "function_name": function_name,
                        "arguments": arguments
                    }
        return None
    
    def _replace_result_references(self, step_description: str, steps: List[Dict[str, Any]]) -> str:
        def replace_match(match):
            step_index = int(match.group(1)) - 1  # 转换为0-based索引
            if 0 <= step_index < len(steps) and steps[step_index]["result"] is not None:
                key = f"[步骤{step_index+1}结果]"
                self.results[key] = steps[step_index]["result"]
                return str(steps[step_index]["result"])
            return match.group(0)  # 如果结果不存在，保持原样
        
        # 修复：使用正则表达式提取键中的数字序号，确保排序正确
        # 原逻辑中的x[2:-3]在中文环境下可能错误截取，改为从键中提取数字部分
        sorted_keys = sorted(self.results.keys(), key=lambda x: int(re.search(r'\d+', x).group()), reverse=True)
        
        for key in sorted_keys:
            if key in step_description:
                step_description = step_description.replace(key, str(self.results[key]))
        
        return self.result_ref_pattern.sub(replace_match, step_description)
    
    def _determine_function_for_step(self, step_description: str, category: str, steps: List[Dict[str, Any]]) -> Dict[str, Any]:
        """为每个步骤确定要调用的函数"""
        # 优先处理简单表达式
        function_info = self._parse_simple_expression(step_description, steps)
        if function_info:
            return function_info
            
        # 根据类别选择工具
        if category == "基础数学":
            tools = self.base_tool["tools"]
        elif category == "三角数学":
            tools = self.trig_tool["tools"]
        elif category == "高等数学":
            tools = self.adv_tool["tools"]
        else:
            tools = self.base_tool["tools"]  # 默认使用基础数学
        
        # 替换步骤描述中的结果引用
        resolved_step_description = self._replace_result_references(step_description, steps)
        
        function_call_prompt = f"""
        请分析以下计算步骤并选择合适的函数进行计算：
        步骤: {resolved_step_description}
        
        可用函数:
        {tools}
        
        请返回函数调用的JSON格式，例如:
        {{
            "name": "函数名",
            "arguments": {{
                "参数1": 值,
                "参数2": 值
            }}
        }}
        """
        
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": function_call_prompt}],
            tools=tools
        )
        
        # 检查是否有函数调用
        if response.choices[0].message.tool_calls:
            tool_call = response.choices[0].message.tool_calls[0]
            function_name = tool_call.function.name
            arguments = tool_call.function.arguments
            
            return {
                "function_name": function_name,
                "arguments": arguments
            }
        else:
            return None
    
    def handle_math_question(self, question: str) -> str:
        """处理数学问题"""
        # 清空之前的步骤结果
        self.results = {}
        
        # 1. 首先询问大模型，看是否能直接回答
        prompt = f"""
        请判断以下数学问题是否可以在您的知识库中直接回答：
        问题: {question}
        
        如果可以，请直接给出答案；如果不能，请返回"需要计算",不要返回额外信息。
        """
        
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.1
        )
        
        answer = response.choices[0].message.content.strip()
        print(f"大模型判断: {answer}")
        
        if answer != "需要计算":
            return answer
        
        # 2. 如果需要计算，确定数学问题类别
        category = self._determine_math_category(question)
        print(f"确定数学问题类别为: {category}")
        
        # 3. 分解问题为多个步骤
        steps = self._decompose_steps(question)
        print(f"分解为 {len(steps)} 个计算步骤")
        
        # 4. 顺序执行每个步骤
        for i, step in enumerate(steps, 1):
            print(f"\n步骤 {i}: {step['description']};")
            
            # 确定要调用的函数，传入已完成的步骤用于引用替换
            function_info = self._determine_function_for_step(step['description'], category, steps[:i-1])
            
            if function_info:
                function_name = function_info["function_name"]
                arguments = function_info["arguments"]
                
                # 执行函数调用
                result = self._execute_function_call(function_name, arguments)
                
                # 存储结果到步骤和results字典
                step["result"] = result
                key = f"[步骤{i}结果]"
                self.results[key] = result
                
                # 添加到对话历史
                self.conversation_history.append({
                    "role": "assistant",
                    "content": f"步骤 {i} 结果: {result}"
                })
                
                print(f"步骤 {i} 执行完成，结果: {result}")
            else:
                step["result"] = "无法确定如何计算该步骤"
                print("无法确定如何计算该步骤")
        
        # 5. 汇总所有步骤的结果
        final_result = "\n".join([
            f"步骤 {i+1}: {step['description']} = {step['result']}"
            for i, step in enumerate(steps)
        ])
        
        return f"计算过程:\n{final_result}\n最终结果: {steps[-1]['result']}"

# 使用示例
if __name__ == "__main__":
    agent = MathAgent()
    
    # 测试问题
    test_questions = [
        "计算 123456 加上  99999999",
        "计算 66666 与 123456 这两个数的乘积与和的差是多少？"
    ]
    
    for question in test_questions:
        print(f"\n问题: {question}")
        result = agent.handle_math_question(question)
        print(f"结果: {result}\n")
